/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 */
#include "typical_qp_manager.h"
#include <map>
#include <memory>
#include "externalinput.h"
#include "adapter_rts_common.h"
#include "adapter_hccp_common.h"
#include "network_manager_pub.h"
#include "rdma_resource_manager.h"

namespace hccl {
TypicalQpManager::TypicalQpManager()
{
}

TypicalQpManager::~TypicalQpManager()
{
    std::unique_lock<std::mutex> lock(qpMutex_);
    for (auto& item : qpMap_) {
        if (item.second.second == nullptr) {
            continue;
        }
        HrtRaQpDestroy(item.second.second);
    }
    qpMap_.clear();
    lock.unlock();
}

TypicalQpManager& TypicalQpManager::GetInstance()
{
    static TypicalQpManager qpInstance;
    return qpInstance;
}

HcclResult TypicalQpManager::CreateQp(struct typical_qp& qpInfo)
{
    HcclResult ret = HCCL_SUCCESS;
    QpHandle qpHandle = nullptr;
    if (!rdmaHandle_) {
        CHK_RET(RdmaResourceManager::GetInstance().GetRdmaHandle(rdmaHandle_));
    }
    CHK_PTR_NULL(rdmaHandle_);
    std::unique_lock<std::mutex> lock(qpMutex_);
    ret = hrtRaTypicalQpCreate(rdmaHandle_, QP_FLAG_RC, OPBASE_QP_MODE, &qpInfo, qpHandle);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[TypicalQpManager][CreateQp] Create qp failed."), HCCL_E_INTERNAL);

    qpMap_.insert(std::make_pair(qpInfo.qpn, std::make_pair(qpInfo, qpHandle)));
    lock.unlock();

    return HCCL_SUCCESS;
}

HcclResult TypicalQpManager::ModifyQp(struct typical_qp& localQpInfo, struct typical_qp& remoteQpInfo)
{
    CHK_PRT_RET((localQpInfo.qpn == 0 || remoteQpInfo.qpn == 0),
        HCCL_ERROR("[TypicalQpManager][ModifyQp] the qpinfo is wrong, qpn is 0."), HCCL_E_PARA);
    QpHandle qpHandle = nullptr;
    CHK_RET(GetQpHandleByQpn(localQpInfo.qpn, qpHandle));
    CHK_PTR_NULL(qpHandle);
    CHK_RET(SetQpRdmaRetryCfg(localQpInfo));
    std::unique_lock<std::mutex> lock(qpMutex_);
    HcclResult ret = hrtRaTypicalQpModify(qpHandle, &localQpInfo, &remoteQpInfo);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[TypicalQpManager][ModifyQp] Modify qp failed."), HCCL_E_INTERNAL);

    return HCCL_SUCCESS;
}

HcclResult TypicalQpManager::DestroyQp(struct typical_qp& qpInfo)
{
    CHK_PRT_RET((qpInfo.qpn == 0), HCCL_ERROR("[TypicalQpManager][DestroyQp] The qpinfo is wrong, qpn is 0."),
        HCCL_E_PARA);
    QpHandle qpHandle;
    CHK_RET(GetQpHandleByQpn(qpInfo.qpn, qpHandle));
    CHK_PTR_NULL(qpHandle);
    std::unique_lock<std::mutex> lock(qpMutex_);
    HcclResult ret = HrtRaQpDestroy(qpHandle);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[TypicalQpManager][DestroyQp] Destory qp failed."), HCCL_E_INTERNAL);
    qpMap_.erase(qpInfo.qpn);
    lock.unlock();
    return HCCL_SUCCESS;
}

HcclResult TypicalQpManager::SetQpRdmaRetryCfg(struct typical_qp& qpInfo)
{
    qpInfo.retry_cnt = GetExternalInputRdmaRetryCnt();
    qpInfo.retry_time = GetExternalInputRdmaTimeOut();
    HCCL_INFO("[TypicalQpManager][SetQpCreateBaseInfo] Qpinfo is setted, tc is %u, sl is %u, retry cnt is %u, "\
        "retry time is %u", qpInfo.tc, qpInfo.sl, qpInfo.retry_cnt, qpInfo.retry_time);
    return HCCL_SUCCESS;
}

HcclResult TypicalQpManager::GetQpHandleByQpn(u32 qpn, QpHandle& qpHandle)
{
    HCCL_DEBUG("[TypicalQpManager][GetQpHandleByQpn] Get qpHandle by qpn[%u]", qpn);
    std::unique_lock<std::mutex> lock(qpMutex_);
    auto it = qpMap_.find(qpn);
    CHK_PRT_RET((it == qpMap_.end()),
        HCCL_ERROR("[TypicalQpManager][GetQpHandleByQpn] Qpn is not found"), HCCL_E_NOT_FOUND);
    qpHandle = it->second.second;
    CHK_PRT_RET((qpHandle == nullptr),
        HCCL_ERROR("[TypicalQpManager][GetQpHandleByQpn] Get Qphandle failed, qphandle is nullptr. qpn is %u", qpn),
        HCCL_E_NOT_FOUND);
    lock.unlock();
    return HCCL_SUCCESS;
}
} // namespace hccl
